﻿//#define NOIMPLICIT
#define DIRECT_EXPLICIT
using System;
using System.Diagnostics;
using RayDen.Library.Components.Surface;
using RayDen.Library.Core;
using RayDen.Library.Core.Primitives;
using RayDen.RayEngine.Core.Interface;
using RayDen.RayEngine.Core.Types;
using RayDen.RayEngine.Engines.PathTracer;
using RayDen.RayEngine.Samplers;

namespace RayDen.RayEngine.Engines.BDPT
{
    public class SBDPTPath : PathSamplerBase
    {
        internal RgbSpectrum Throughput, LightThroughtput;
        public PathTracerPathState PathState;


        protected int lightDepth, depth, tracedShadowRayCount, currentLightVertex;
        protected float pathWeight;
        protected bool specularBounce;

        public UInt64 PathExpression;

        protected RayEngineScene scene;

        private bool lightStoped;
        private RayData LightRay;
        private int LightRayIndex;

        private ShadowRayInfo[] secRays;
        private LightVertex[] lightVertices;
        
        SurfaceIntersectionData hitInfo = null, lightHitInfo = null;
        private float lightWeight;
        public int MaxRaysPerPath
        {
            get;
            set;
        }

        public struct LightVertex
        {
            public Point HitPoint;
            public Normal Normal;
            public RgbSpectrum Throughput;
            public Vector Wi;
            public RgbSpectrum BsdfSample;
            public float BsdfWeight;
            public float dAWeight;
        }

        private void RestartLight()
        {
            this.LightRayIndex = -1;
            float lPdf;
            LightSample lightSample;
            var light = scene.Lights[scene.SampleLights(this.Sample.GetLazyValue(), out lPdf)];
            light.EvaluatePhoton(scene, this.Sample.GetLazyValue(), this.Sample.GetLazyValue(), this.Sample.GetLazyValue(), this.Sample.GetLazyValue(),
                                 this.Sample.GetLazyValue(),
                                 out lightSample);
            LightThroughtput = (RgbSpectrum)(RgbSpectrumInfo)(lightSample.Spectrum) / (lightSample.Pdf);
            lightWeight = lightSample.Pdf*lPdf;
            LightRay = lightSample.LightRay;
            this.currentLightVertex = 0;
            this.lightDepth = 0;

        }


        public override void InitPath(IPathProcessor buffer)
        {
            base.InitPath(buffer);
            this.scene = pathIntegrator.Scene;
            this.Radiance = new RgbSpectrum(0f);
            this.Throughput = new RgbSpectrum(1f);
            this.PathState = PathTracerPathState.EyeVertex;
            if (this.lightVertices == null)
                this.lightVertices = new LightVertex[scene.MaxPathDepth];
            if (this.secRays == null)
                this.secRays = new ShadowRayInfo[scene.ShadowRaysPerSample + scene.MaxPathDepth];
            this.Sample = pathIntegrator.Sampler.GetSample(this.Sample ?? null);
            pathIntegrator.Scene.Camera.GetRay(Sample.imageX, Sample.imageY, out this.PathRay);
            this.RayIndex = -1;
            this.pathWeight = 1.0f;
            this.tracedShadowRayCount = 0;
            this.depth = 0;
            this.specularBounce = true;
            this.RestartLight();
        }

        public override bool FillRayBuffer(Core.Types.RayBuffer rayBuffer)
        {
            var leftSpace = rayBuffer.LeftSpace();
            if (((PathState == PathTracerPathState.EyeVertex) && (2 > leftSpace)) ||
            ((PathState == PathTracerPathState.ShadowRaysOnly) && (tracedShadowRayCount > leftSpace)) ||
            ((PathState == PathTracerPathState.NextVertex) && (tracedShadowRayCount + 2 > leftSpace)))
                return false;

            if (PathState != PathTracerPathState.ShadowRaysOnly)
            {
                RayIndex = rayBuffer.AddRay(ref PathRay);
                if (!lightStoped)
                    LightRayIndex = rayBuffer.AddRay(ref LightRay);
            }

            if (PathState == PathTracerPathState.NextVertex || PathState == PathTracerPathState.ShadowRaysOnly)
            {
                for (int i = 0; i < tracedShadowRayCount; ++i)
                    secRays[i].ShadowRayIndex = rayBuffer.AddRay(ref secRays[i].ShadowRay);
            }
            return true;
        }

        public override void Advance(Core.Types.RayBuffer rayBuffer, Core.Types.SampleBuffer consumer)
        {
            int currentTriangleIndex = 0;
#if VERBOSE

            try {
#endif

            if (((PathState == PathTracerPathState.ShadowRaysOnly) || (PathState == PathTracerPathState.NextVertex)) &&
                (tracedShadowRayCount > 0))
            {
                for (int i = 0; i < tracedShadowRayCount; ++i)
                {
                    RayHit shadowRayHit = rayBuffer.rayHits[secRays[i].ShadowRayIndex];
                    RgbSpectrum attenuation;
                    bool continueTrace;
                    if (this.ShadowRayTest(ref shadowRayHit, ref secRays[i].ShadowRay, out attenuation, out continueTrace))
                    {
                        float weight = 1f;//tracedShadowRayCount;
                        //                            Radiance.MADD()
                        Radiance += attenuation * (secRays[i].Throughput /(weight*secRays[i].Pdf))*100f;
                        pathWeight *= secRays[i].Pdf;
                    }
                }

                if (PathState == PathTracerPathState.ShadowRaysOnly)
                {
                    Splat(consumer);
                    return;
                }
                tracedShadowRayCount = 0;
            }
            RayHit eyeRayHit = rayBuffer.rayHits[RayIndex];
            Vector wo = -PathRay.Dir;

            depth++;
            bool eyeMissed = eyeRayHit.Index == 0xffffffffu;
            if (eyeMissed && lightDepth > 0 && tracedShadowRayCount > 0)
            {
                Console.WriteLine("!");
                
            }

            if (eyeMissed || PathState == PathTracerPathState.ShadowRaysOnly || depth > scene.MaxPathDepth)
            {
                if (eyeMissed)
                {
                    Radiance += this.scene.SampleEnvironment(ref wo) * Throughput;
                }
                Splat(consumer);

                return;
            }

            // Something was hit
            if (hitInfo == null)
            {
                hitInfo = SurfaceSampler.GetIntersection(ref PathRay, ref eyeRayHit);
            }
            else
            {
                SurfaceSampler.GetIntersection(ref PathRay, ref eyeRayHit, ref hitInfo);
            }
            currentTriangleIndex = (int)eyeRayHit.Index;

            if (hitInfo == null)
            {
                Debugger.Break();
            }
            //If Hit light)

            if (hitInfo.IsLight)
            {
                if (specularBounce || depth == 1)
                {
                    var lt = scene.GetLightByIndex(currentTriangleIndex);
                    if (lt != null)
                    {
#if !NOIMPLICIT
                        var le = (RgbSpectrumInfo)(lt.Le(ref wo));
                        Radiance += Throughput * (RgbSpectrum)le;
#endif
                    }
                }
                Splat(consumer);
                return;
            }

            if (!lightStoped)
            {
                RayHit lightRayHit = rayBuffer.rayHits[LightRayIndex];
                if (lightRayHit.Miss())
                {
                    this.RestartLight();
                }
                else
                {
                    lightDepth++;
                    if (lightHitInfo == null)
                    {
                        lightHitInfo = SurfaceSampler.GetIntersection(ref LightRay, ref lightRayHit,
                                                                      IntersectionOptions.ResolveSurfaceComplete |
                                                                      IntersectionOptions.ResumeOnLight);
                    }
                    else
                    {
                        SurfaceSampler.GetIntersection(ref LightRay, ref lightRayHit, ref lightHitInfo,
                                                                      IntersectionOptions.ResolveSurfaceComplete |
                                                                      IntersectionOptions.ResumeOnLight);
                    }

                    float lPdf = 0f;
                    var lwi = new Vector();
                    var lhp = LightRay.Point(lightRayHit.Distance);
                    var lf = lightHitInfo.MMaterial.Sample_f(ref wo, out lwi, ref hitInfo.Normal,
                                                             ref hitInfo.ShadingNormal, ref Throughput,
                                                             Sample.GetLazyValue(), Sample.GetLazyValue(),
                                                             Sample.GetLazyValue(), ref hitInfo.TextureData,
                                                             out lPdf, out specularBounce);
                    var scatteringCorrection = Vector.AbsDot(ref wo, ref hitInfo.Normal) /
                                                    Vector.AbsDot(ref wo, ref hitInfo.ShadingNormal);
                    scatteringCorrection *=  Vector.AbsDot(ref lwi, ref hitInfo.ShadingNormal);
                    if (lPdf <= 0f || lf.IsBlack())
                    {
                        this.RestartLight();
                        goto eyeCont;
                    }
                    if (lightDepth 
#if DIRECT_EXPLICIT
                        > 1)
#else
                        > 0)
#endif

                    {

                        float da = (currentLightVertex > 0 ? this.lightVertices[currentLightVertex - 1].dAWeight : 1.0f)*lightWeight;

                        this.lightVertices[currentLightVertex] = new LightVertex()
                            {
                                HitPoint = lhp,
                                Normal = lightHitInfo.ShadingNormal,
                                BsdfWeight = lPdf,
                                BsdfSample = lf*scatteringCorrection,
                                dAWeight = da*(currentLightVertex > 0 
                                            ?
                                                lPdf*Geometry.G(ref lhp, ref lightVertices[currentLightVertex-1].HitPoint,ref hitInfo.ShadingNormal,ref lightVertices[currentLightVertex-1].Normal)
                                            : lPdf)*scatteringCorrection,
                                Throughput = LightThroughtput,
                                Wi = LightRay.Dir,
                            };
                        currentLightVertex++;
                    }
                    
                    LightThroughtput *= (lf / lPdf)*scatteringCorrection;
                    LightRay = new RayData(ref lhp, ref lwi);
                }
                lightStoped = currentLightVertex > scene.MaxPathDepth || LightThroughtput.IsBlack();
            }

            //if (currentLightVertex > 3)
            //{
            //    Debugger.Break();
            //}
        eyeCont:
            var hitPoint = PathRay.Point(eyeRayHit.Distance);

            tracedShadowRayCount = 0;
            var bsdf = hitInfo.MMaterial;

            if (bsdf.IsDiffuse())
            {
#if DIRECT_EXPLICIT
                float lightStrategyPdf = LightSampler.StrategyPdf;
                //scene.ShadowRaysPerSample/(float)scene.Lights.Length;
                RgbSpectrum lightTroughtput = Throughput * hitInfo.Color;
                LightSampler.EvaluateShadow(ref hitPoint, ref hitInfo.Normal, Sample.GetLazyValue(),
                                            Sample.GetLazyValue(), Sample.GetLazyValue(), ref ls);
                for (int index = 0; index < ls.Length; index++)
                {

                    if (ls[index].Pdf <= 0f)
                        continue;
                    secRays[tracedShadowRayCount].Throughput = (RgbSpectrum)(RgbSpectrumInfo)(ls[index].Spectrum);
                    secRays[tracedShadowRayCount].Pdf = ls[index].Pdf;
                    secRays[tracedShadowRayCount].ShadowRay = ls[index].LightRay;
                    Vector lwi = secRays[tracedShadowRayCount].ShadowRay.Dir;
                    RgbSpectrum fs;
                    hitInfo.MMaterial.f(ref secRays[tracedShadowRayCount].ShadowRay.Dir, ref wo, ref hitInfo.ShadingNormal, ref lightTroughtput, out fs, BrdfType.Diffuse);
                    secRays[tracedShadowRayCount].Throughput *= lightTroughtput *
                                                                Vector.AbsDot(ref hitInfo.Normal, ref lwi) *
                                                                fs;
                    if (!secRays[tracedShadowRayCount].Throughput.IsBlack())
                    {
                        secRays[tracedShadowRayCount].Pdf /= lightStrategyPdf;
                        tracedShadowRayCount++;
                    }
                }
#endif
                for (int index = 0; index < currentLightVertex; index++)
                {
                    if (lightVertices[index].Throughput.IsBlack())
                        continue;

                    var lwi = lightVertices[index].Wi;
                    RgbSpectrum fs;
                    hitInfo.MMaterial.f(ref lwi, ref wo, ref hitInfo.ShadingNormal, ref Throughput, out fs, types: BrdfType.Diffuse);
                    var dir = lightVertices[index].HitPoint - hitPoint;
                    var dirLength = dir.Length;
                    dir.NormalizeSelf();

                    secRays[tracedShadowRayCount].Throughput = Throughput*lightVertices[index].Throughput
                                                               *fs
                                                               *lightVertices[index].BsdfSample
                                                               //* Vector.AbsDot(ref hitInfo.Normal, ref lwi) 
                                                               *Geometry.G(ref hitPoint, ref lightVertices[index].HitPoint,ref hitInfo.Normal,ref lightVertices[index].Normal);

                    secRays[tracedShadowRayCount].Pdf = (pathWeight*lightVertices[index].dAWeight)
                        ;
                    //    * Geometry.G(ref hitPoint, ref lightVertices[index].HitPoint, ref hitInfo.Normal, ref lightVertices[index].Normal);
                        //1f/(depth+index) * (pathWeight*lightVertices[index].dAWeight);
                    secRays[tracedShadowRayCount].ShadowRay = new RayData(ref hitPoint, ref dir, 1e-4f, dirLength-1e-4f);
                    tracedShadowRayCount++;
                }
            }

            float fPdf = 0f;
            var wi = new Vector();
            RgbSpectrum f;

            f = bsdf.Sample_f(ref wo, out wi, ref hitInfo.Normal, ref hitInfo.ShadingNormal, ref Throughput,
                                Sample.GetLazyValue(), Sample.GetLazyValue(),
                                Sample.GetLazyValue(), ref hitInfo.TextureData,
                                out fPdf, out specularBounce);

            if ((fPdf <= 0.0f) || f.IsBlack())
            {
                if (tracedShadowRayCount > 0)
                    PathState = PathTracerPathState.ShadowRaysOnly;
                else
                {
                    Splat(consumer);

                }
                return;
            }
            pathWeight *= fPdf;
            Throughput *= (f) / fPdf;

            if (depth > scene.MaxPathDepth)
            {
                float prob = Math.Max(Throughput.Filter(), scene.RussianRuletteImportanceCap);
                if (prob >= Sample.GetLazyValue())
                {
                    Throughput /= prob;
                    pathWeight *= prob;
                }
                else
                {
                    if (tracedShadowRayCount > 0)
                        PathState = PathTracerPathState.ShadowRaysOnly;
                    else
                    {
                        Splat(consumer);

                    }

                    return;
                }
            }

            PathRay.Org = hitPoint;
            PathRay.Dir = wi.Normalize();
            PathState = PathTracerPathState.NextVertex;

#if VERBOSE
            }
            catch (Exception ex) {
                RayDen.Library.Components.SystemComponents.Tracer.TraceLine("Advance path exception");
                RayDen.Library.Components.SystemComponents.Tracer.TraceLine("Error triangle {0}", currentTriangleIndex);
                RayDen.Library.Components.SystemComponents.Tracer.TraceLine(ex.Message);
                RayDen.Library.Components.SystemComponents.Tracer.TraceLine(ex.StackTrace);

            }
#endif
        }


    }
}
